/*
 * sparse_vec.cpp
 *
 *  Created on: Feb 28, 2011
 *      Author: tqlong
 */

#include "sparse_vec.h"
#include <cstdio>
#include <cmath>
#include <armadillo>
#include <cassert>
#include <iostream>

ARMA_NAMESPACE_BEGIN;

const sparse_vec& sparse_vec::randn()
{
  nz_.clear();
  idx_.clear();
  for (int i = 0; i < dim_; i++) {
//    std::cerr << "-- rand i = " << i << "\n";
    double r = arma::randu(1)[0];
    if (r < 0.5) {                                     // 0.5 chance to be non-zero
      double v = arma::randn(1)[0];                    // Normal distribution
      nz_.push_back(v);
      idx_.push_back(i);
//      std::cerr << "rand i = " << i << " val = " << v << "\n";
    }
  }
  return *this;
}

const sparse_vec& sparse_vec::randu()
{
  nz_.clear();
  idx_.clear();
  for (int i = 0; i < dim_; i++) {
    double r = arma::randu(1)[0];
    if (r < 0.5) {                                     // 0.5 chance to be non-zero
      double v = arma::randu(1)[0];                    // Uniform distribution
      nz_.push_back(v);
      idx_.push_back(i);
    }
  }
  return *this;
}

const sparse_vec& sparse_vec::zeros()
{
  nz_.clear();                                         // clear all non-zeros
  idx_.clear();
  return *this;
}

const sparse_vec& sparse_vec::ones()
{
  nz_ = std::vector<double>(dim_, 1.0);                // all non-zeros are 1
  idx_ = std::vector<int>(dim_);
  for (int i = 0; i < dim_; i++)
    idx_[i] = i;
  return *this;
}

double sparse_vec::get(int i) const
{
//  std::cerr<< "const\n";
  int pos;
  if (!find(i, pos)) return 0.0;                       // if zero
  else return nz_[pos];
}

double& sparse_vec::ref(int i)
{
//  std::cerr<< "ref\n";
  int pos;
  if (!find(i, pos)) {                                 // if zero
    nz_.insert(nz_.begin()+pos, 0.0);                  // add a new non-zero position
    idx_.insert(idx_.begin()+pos, i);
  }
  return nz_[pos];
}

bool sparse_vec::find(int i, int &pos) const
{
  bool ret = false;
  pos = nnz();
  if (nnz() == 0) return ret;

  int l = 0, r = nnz()-1;
//  std::cerr << "i = " << i << "\n";
  while (l <= r) {
    int c = (l+r)/2;
//    std::cerr << l << " " << c << " = " << idx_[c] << " " << r << "\n";
    if (idx_[c] < i) l = c+1;
    else if (idx_[c] > i) r = c-1;
    else {
      pos = c;
      ret = true;
      break;
    }
  }
  if (ret) return ret;
  else {  // r = l-1
    if (r < 0) pos = 0;
    else if (idx_[r] < i) pos = r+1;
    else pos = r;
  }

//  std::cerr << "i = " << i << " pos = " << pos << " ret = " << ret << "\n";
  return ret;
}

sparse_vec sparse_vec::operator/(double s) const
{
  return (*this)*(1.0/s);
}

sparse_vec sparse_vec::operator*(double s) const
{
  sparse_vec v(*this);
  for (int i = 0; i < nnz(); i++)                      // all nonzeros is scaled by s
    v.nz_[i] *= s;
  return v;
}

sparse_vec sparse_vec::operator-() const
{
  sparse_vec v(*this);
  for (int i = 0; i < nnz(); i++)
    v.nz_[i] = -v.nz_[i];
  return v;
}

sparse_vec sparse_vec::operator+() const
{
  return (*this);
}

const sparse_vec& sparse_vec::operator*=(double s)
{
  for (int i = 0; i < nnz(); i++)                      // all nonzeros is scaled by s
    nz_[i] *= s;
  return *this;
}

const sparse_vec& sparse_vec::operator/=(double s)
{
  return (*this*=(1.0/s));
}

sparse_vec sparse_vec::operator-(const sparse_vec& s) const
{
  assert(dim_ == s.dim_);
  sparse_vec v(dim_);
  int i = 0, j = 0;
  while (i < nnz() || j < s.nnz()) {
    int uid = i < nnz() ? idx_[i] : dim_;              // check current pointer
    int vid = j < s.nnz() ? s.idx_[j] : dim_;
//    std::cout << "i = " << i << " j = " << j << " uid = " << uid << " vid = " << vid << "\n";
    if (uid < vid) {
      v.nz_.push_back(nz_[i]);                         // if u and !v
      v.idx_.push_back(uid);
      i++;
    }
    else if (uid > vid) {                              // if !u and v
      v.nz_.push_back(-s.nz_[j]);
      v.idx_.push_back(vid);
      j++;
    }
    else {                                             // if u and v
      v.nz_.push_back(nz_[i]-s.nz_[j]);
      v.idx_.push_back(uid);
      i++;
      j++;
    }
  }
  return v;
}

sparse_vec sparse_vec::operator+(const sparse_vec& s) const
{
  assert(dim_ == s.dim_);
  sparse_vec v(dim_);
  int i = 0, j = 0;
  while (i < nnz() || j < s.nnz()) {
    int uid = i < nnz() ? idx_[i] : dim_;              // check current pointer
    int vid = j < s.nnz() ? s.idx_[j] : dim_;
//    std::cout << "i = " << i << " j = " << j << " uid = " << uid << " vid = " << vid << "\n";
    if (uid < vid) {
      v.nz_.push_back(nz_[i]);                         // if u and !v
      v.idx_.push_back(uid);
      i++;
    }
    else if (uid > vid) {                              // if !u and v
      v.nz_.push_back(s.nz_[j]);
      v.idx_.push_back(vid);
      j++;
    }
    else {                                             // if u and v
      v.nz_.push_back(nz_[i]+s.nz_[j]);
      v.idx_.push_back(uid);
      i++;
      j++;
    }
  }
  return v;
}

const sparse_vec& sparse_vec::operator+=(const sparse_vec& s)
{
//  assert(dim_ == s.dim_);
//  sparse_vec v(dim_);
//  int i = 0, j = 0;
//  while (i < nnz() || j < s.nnz()) {                   // similar to operator+-
//    int uid = i < nnz() ? idx_[i] : dim_;
//    int vid = j < s.nnz() ? s.idx_[j] : dim_;
//    if (uid < vid) {
//      v.nz_.push_back(nz_[i]);
//      v.idx_.push_back(uid);
//      i++;
//    }
//    else if (uid > vid) {
//      v.nz_.push_back(s.nz_[j]);
//      v.idx_.push_back(vid);
//      j++;
//    }
//    else {
//      v.nz_.push_back(nz_[i]+s.nz_[j]);
//      v.idx_.push_back(uid);
//      i++;
//      j++;
//    }
//  }
//  *this = v;
  *this = *this + s;
  return *this;
}

const sparse_vec& sparse_vec::operator-=(const sparse_vec& s)
{
//  assert(dim_ == s.dim_);
//  sparse_vec v(dim_);
//  int i = 0, j = 0;
//  while (i < nnz() || j < s.nnz()) {                   // similar to operator+-
//    int uid = i < nnz() ? idx_[i] : dim_;
//    int vid = j < s.nnz() ? s.idx_[j] : dim_;
//    if (uid < vid) {
//      v.nz_.push_back(nz_[i]);
//      v.idx_.push_back(uid);
//      i++;
//    }
//    else if (uid > vid) {
//      v.nz_.push_back(-s.nz_[j]);
//      v.idx_.push_back(vid);
//      j++;
//    }
//    else {
//      v.nz_.push_back(nz_[i]-s.nz_[j]);
//      v.idx_.push_back(uid);
//      i++;
//      j++;
//    }
//  }
//  *this = v;
  *this = *this - s;
  return *this;
}

double norm(const sparse_vec& v, int ell)
{
  double s = 0;
  switch (ell) {
  case 1:
    for (int i = 0; i < v.nnz(); i++)                  // l1 = sum_|x_i|
      s += ::fabs(v.nz_[i]);
    break;
  case 2:
    for (int i = 0; i < v.nnz(); i++)                  // l2 = sqrt(sum x_i^2)
      s += v.nz_[i]*v.nz_[i];
    s = ::sqrt(s);
    break;
  default:
    for (int i = 0; i < v.nnz(); i++)                  // l-inf = max |x_i|
      if (s < ::fabs(v.nz_[i])) s = ::fabs(v.nz_[i]);
    break;
  }
  return s;
}

double norm(const sparse_vec& v, const char* type)
{
  if (strcmp(type, "1") == 0)
    return norm(v, 1);
  else if (strcmp(type, "2") == 0)
    return norm(v, 2);
  else if (strcmp(type, "inf") == 0)
    return norm(v, -1);
  return 0.0;
}

double dot(const sparse_vec& u, const sparse_vec& v)
{
  int i = 0, j = 0;
  double s = 0;
  while (i < u.nnz() && j < v.nnz()) {
    if (u.idx_[i] < v.idx_[j]) i++;
    else if (u.idx_[i] > v.idx_[j]) j++;
    else {
      s += u.nz_[i]*v.nz_[j];                          // only take product of same index pair
      i++;
      j++;
    }
  }
  return s;
}

ARMA_NAMESPACE_END;

std::ostream& operator<<(std::ostream& os, const arma::sparse_vec& v)
{
  for (int i = 0; i < v.nnz(); i++)
    os << v.idx()[i] << ":" << v.nz()[i] << " ";
  return os;
}

arma::sparse_vec operator*(double s, const arma::sparse_vec& v)
{
  return v*s;
}
